// Accelerated Dual Decomposition - An implementation of a MAP inference
// algorithm. See ai.stanford.edu/~vjojic/fastdd.pd
//
// Copyright 2011 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: nadavs@google.com (Nadav Samet)
//          noam.levy1@stanford.edu (Noam Levy)
//
// Unit tests for AcceleratedDD.

#include "AcceleratedDD.h"

#include <iostream>
#include <assert.h>
#include <stdlib.h>

using namespace AcceleratedDD;

bool TestIncrementAssignment() {
  size_t ranges[5] = {3,3,2,4,5};
  vector<VariableId> scope;
  scope.push_back(1);
  scope.push_back(2);
  scope.push_back(4);
  size_t assignment[5] = {0, 0, 0, 0, 0};

  assert(IncrementAssignment(assignment, ranges, scope));
  assert(assignment[1]==1);
  assert(assignment[2]==0);
  assert(assignment[4]==0);

  assert(IncrementAssignment(assignment, ranges, scope));
  assert(IncrementAssignment(assignment, ranges, scope));
  assert(assignment[1]==0);
  assert(assignment[2]==1);
  assert(assignment[4]==0);

  assignment[1]=2;
  assignment[2]=1;
  assignment[4]=3;
  assert(IncrementAssignment(assignment, ranges, scope));
  assert(assignment[1]==0);
  assert(assignment[2]==0);
  assert(assignment[4]==4);

  assignment[1]=2;
  assignment[2]=1;
  assignment[4]=4;
  assert(!IncrementAssignment(assignment, ranges, scope));
  assert(assignment[1]==0);
  assert(assignment[2]==0);
  assert(assignment[4]==0);
  return true;
}

bool TestUpdateP() {
  OptimizationProblem p;
  p.variables = 9;
  p.variable_range = new size_t[9];
  for (int i = 0;i < 9; ++i) {
    p.variable_range[i] = 4;
  };
  
  for (int i = 0; i < 8; ++i) {
    Scope s;
    s.push_back(i);
    s.push_back(i+1);
    TableFactor *f = new TableFactor(s, &p);
    size_t pair_array[2];
    for (int j = 0; j < 4; ++j) {
      pair_array[0] = pair_array[1] = j;
      f->Assign(pair_array, 1.0);
    }
    f->Assign(pair_array, 2.0);
    p.factors.push_back(f);
  }

  // singleton
  for (int i = 0; i < 9; ++i) {
    Scope s;
    s.push_back(i);
    TableFactor *f = new TableFactor(s, &p);
    size_t array[1];
    for (int j = 0; j < 4; ++j) {
      array[0] = j;
      f->Assign(array, j);
    }
    p.factors.push_back(f);
  }

  AlgorithmState *s = CreateInitialAlgorithmState(p);
  Iterate(p, 1.0, 1.0, s);
  cerr << "Dual value " << s->dual_value << endl;
  return true;
}

int main(int argc, char **argv) {
  if (!TestIncrementAssignment()) {
    cerr << "TestIncrementAssignment failed.";
  }

  if (!TestUpdateP()) {
    cerr << "TestUpdateP failed.";
  }
}
